from collections.abc import Generator
import time

from assertpy import assert_that
from kubernetes.client import V1Pod
from playwright.sync_api import Page, BrowserContext
import pytest

from utils.client import KubeClient, StressConfig
import actions.colocation.colocation_configuration as cc


THRESHOLD_LOW = 10
THRESHOLD_HIGH = 90
STRESS_LOAD = 60


@pytest.fixture(scope="function")
def conf_page(context: BrowserContext) -> Generator[Page, None, None]:
    """conf_page默认为启动负载感知调度，但是关闭真实负载调度阈值"""
    page = context.new_page()
    cc.visit_colocation_config_page(page)
    cc.do_open_modal(page)
    if not cc.get_usage_plugin_switch(page).is_checked():
        cc.do_toggle_usage_plugin_switch(page)
    if cc.get_usage_threshold_switch(page).is_checked():
        cc.do_toggle_usage_threshold_switch(page)
    cc.do_click_modal_confirm(page)
    yield page

    cc.do_open_modal(page)
    if cc.get_usage_plugin_switch(page).is_checked():
        cc.do_toggle_usage_plugin_switch(page)
        cc.do_click_modal_confirm(page)
    page.close()


@pytest.fixture(scope="module")
def stress_pod(kube_client: KubeClient) -> Generator[V1Pod, None, None]:
    stress_pod = kube_client.deploy_stress_pod()
    yield stress_pod

    name, ns = stress_pod.metadata.name, stress_pod.metadata.namespace
    kube_client.delete_pod(name, ns, 60)


@pytest.fixture(scope="function")
def stress_up_cpu(kube_client: KubeClient, stress_pod: V1Pod, numa_stat: dict) -> Generator[None, None]:
    kube_client.start_stress_process(
        name=stress_pod.metadata.name,
        namespace=stress_pod.metadata.namespace,
        stress_config=StressConfig(
            cpu_core=numa_stat["cpus"], 
            cpu_load=STRESS_LOAD,
        ),
    )
    time.sleep(10)
    yield

    kube_client.stop_stress_process()


@pytest.fixture(scope="function")
def stress_up_memory(kube_client: KubeClient, stress_pod: V1Pod, numa_stat: dict) -> Generator[None, None]:
    kube_client.start_stress_process(
        name=stress_pod.metadata.name,
        namespace=stress_pod.metadata.namespace,
        stress_config=StressConfig(
            vm=1,
            vm_bytes=STRESS_LOAD,
        ),
    )
    time.sleep(10)
    yield

    kube_client.stop_stress_process()


def update_threshold(page: Page, cpu: float = 0, mem: float = 0, wait_time: int = 0):
    if not cpu or not mem:
        return
    cc.do_open_modal(page)
    if not cc.get_usage_threshold_switch(page).is_checked():
        cc.do_toggle_usage_threshold_switch(page)
    cpu_box, mem_box = cc.get_usage_threshold_textbox(page)
    if cpu >= 0 and cpu <= 100:
        cpu_box.fill(str(cpu))
    if mem >= 0 and mem <= 100:
        mem_box.fill(str(mem))
    cc.do_click_modal_confirm(page)
    time.sleep(wait_time)


class TestUsageThresholdOff:
    """
     @description 测试**真实负载调度**特性关闭时的，不会发生驱逐
    """
    @pytest.mark.usefixtures("stress_up_cpu")
    @staticmethod
    def test_cpu_reaches_threshold(kube_client: KubeClient, conf_page: Page):   # conf_page是需要执行的fixture，不要删
        """Test 77"""
        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            try:
                assert_that(kube_client.wait_pod_condition(name, ns, "Running")).is_true()
            finally:
                kube_client.delete_pod(name, ns)

    @pytest.mark.usefixtures("stress_up_memory")
    @staticmethod
    def test_memory_reaches_threshold(kube_client: KubeClient, conf_page: Page):
        """Test 79"""
        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            try:
                assert_that(kube_client.wait_pod_condition(name, ns, "Running")).is_true()
            finally:
                kube_client.delete_pod(name, ns)


class TestUsageThresholdOn:
    """
     @description 测试**真实负载调度**特性开启时的调度
    """
    @pytest.mark.usefixtures("stress_up_cpu")
    @staticmethod
    def test_cpu_reaches_threshold(kube_client: KubeClient, conf_page: Page):
        """Test 80, 81"""
        update_threshold(conf_page, THRESHOLD_LOW, 100, 10)
        assert_that(kube_client.wait_for_volcano_restart()).is_true()

        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            assert_that(kube_client.wait_pod_condition(name, ns, "Pending")).is_true()
            kube_client.delete_pod(name, ns)

        update_threshold(conf_page, THRESHOLD_HIGH, 100, 10)
        assert_that(kube_client.wait_for_volcano_restart()).is_true()

        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            assert_that(kube_client.wait_pod_condition(name, ns, "Running")).is_true()
            kube_client.delete_pod(name, ns)
        
    @pytest.mark.usefixtures("stress_up_memory")
    @staticmethod
    def test_memory_reaches_threshold(kube_client: KubeClient, conf_page: Page):
        """Test 82, 83"""
        update_threshold(conf_page, 100, THRESHOLD_LOW, 10)
        assert_that(kube_client.wait_for_volcano_restart()).is_true()

        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            assert_that(kube_client.wait_pod_condition(name, ns, "Pending")).is_true()
            kube_client.delete_pod(name, ns)

        update_threshold(conf_page, 100, THRESHOLD_HIGH, 10)
        assert_that(kube_client.wait_for_volcano_restart()).is_true()
        
        for install_func in [kube_client.install_be_pod, kube_client.install_ls_pod, kube_client.install_hls_pod]:
            pod = install_func()
            name, ns = pod.metadata.name, pod.metadata.namespace
            assert_that(kube_client.wait_pod_condition(name, ns, "Running")).is_true()
            kube_client.delete_pod(name, ns)